import pandas as pd
import numpy as np
from scipy import interpolate
import matplotlib.pyplot as plt
from pathlib import Path


def load_and_preprocess_data(file_path, level_col='H', capacity_col='V'):
    """加载并预处理水位库容数据"""
    df = pd.read_excel(file_path)
    if level_col not in df.columns or capacity_col not in df.columns:
        raise ValueError(f"数据列名错误，需要包含'{level_col}'和'{capacity_col}'列")
    df_clean = df[[level_col, capacity_col]].dropna()
    if len(df_clean) < 2:
        raise ValueError("有效数据点不足，至少需要2个非空数据点")
    return df_clean.sort_values(by=level_col).drop_duplicates(subset=level_col)


def perform_interpolation(df, interp_method='linear', num_points=85):
    """执行插值计算"""
    levels = df['H'].values
    capacities = df['V'].values
    interp_func = interpolate.interp1d(levels, capacities, kind=interp_method, fill_value="extrapolate")
    interp_levels = np.linspace(levels.min(), levels.max(), num_points)
    return pd.DataFrame({
        '插值水位': interp_levels,
        '插值库容': interp_func(interp_levels)
    })


def save_results_to_excel(original_df, interp_df, output_path='水位库容插值结果.xlsx'):
    """
    保存原始数据和插值结果到Excel
    :param original_df: 原始数据DataFrame
    :param interp_df: 插值结果DataFrame
    :param output_path: 输出文件路径
    """
    with pd.ExcelWriter(output_path) as writer:
        original_df.to_excel(writer, sheet_name='原始数据', index=False)
        interp_df.to_excel(writer, sheet_name='插值结果', index=False)
    print(f"\n结果已成功保存到：{output_path}")


def visualize_results(original_df, interp_df, method):
    """可视化原始数据和插值结果"""
    plt.figure(figsize=(10, 6))
    plt.scatter(original_df['H'], original_df['V'], color='red', label='原始数据点', zorder=2)
    plt.plot(interp_df['插值水位'], interp_df['插值库容'], color='blue', label=f'{method}插值结果', zorder=1)

    plt.title('水位-库容插值结果对比', fontsize=14)
    plt.xlabel('水位（m）', fontsize=12)
    plt.ylabel('库容（万m³）', fontsize=12)
    plt.legend()
    plt.grid(linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # ================== 配置参数 ==================
    INPUT_EXCEL = '建模资料1.xlsx'  # 输入数据文件路径
    OUTPUT_EXCEL = '水位库容插值结果.xlsx'  # 输出结果文件路径
    LEVEL_COL = 'H'  # 水位列列名
    CAPACITY_COL = 'V'  # 库容列列名
    INTERP_METHOD = 'linear'  # 插值方法（linear/cubic/quadratic）
    # ==============================================

    try:
        # 生成示例数据（若输入文件不存在）
        if not Path(INPUT_EXCEL).exists():
            sample_data = pd.DataFrame({
                LEVEL_COL: [100, 105, 110, 115, 120],
                CAPACITY_COL: [500, 1200, 2500, 4300, 6800]
            })
            sample_data.to_excel(INPUT_EXCEL, index=False)
            print(f"已生成示例输入文件：{INPUT_EXCEL}")

        # 加载并预处理数据
        original_df = load_and_preprocess_data(INPUT_EXCEL, LEVEL_COL, CAPACITY_COL)

        # 执行插值计算
        interp_df = perform_interpolation(original_df, interp_method=INTERP_METHOD)

        # 保存结果到Excel
        save_results_to_excel(original_df, interp_df, OUTPUT_EXCEL)

        # 打印前5个插值结果
        print("\n插值结果示例（前5个点）:")
        print(interp_df.head().round(2))

        # 可视化结果
        visualize_results(original_df, interp_df, INTERP_METHOD)

    except Exception as e:
        print(f"程序运行出错: {str(e)}")
